import time
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict, Set, Any, Optional
from matplotlib.animation import FuncAnimation

from env.grid_env import GridMap


class Simulator:
    """
    仿真引擎，用于执行和可视化机器人路径
    """
    def __init__(self, grid_map: GridMap, config: Dict[str, Any]):
        """
        初始化仿真引擎
        
        Args:
            grid_map: 网格地图
            config: 配置参数
        """
        self.grid_map = grid_map
        self.config = config
        
        # 仿真参数
        self.visualize = config.get('visualize', True)
        self.step_delay = config.get('step_delay', 0.5)  # 增加默认延迟到0.5秒，使多机器人运动更清晰可见
        
        # 仿真状态
        self.current_timestep = 0
        self.robot_positions = {}  # 当前机器人位置 {robot_id: position}
        self.robot_paths = {}  # 机器人路径 {robot_id: [pos1, pos2, ...]}
        self.robot_goals = {}  # 机器人目标 {robot_id: goal}
        self.completed_robots = set()  # 已完成任务的机器人
        
        # 统计信息
        self.conflicts = []  # 冲突列表 [(agent_i, agent_j, pos, timestep), ...]
        self.makespan = 0  # 完工时间
        self.sum_of_costs = 0  # 总代价
    
    def load_solution(self, solution: Dict[int, List[Tuple[int, int]]], goals: Dict[int, Tuple[int, int]]):
        """
        加载路径解决方案
        
        Args:
            solution: 路径字典 {robot_id: [pos1, pos2, ...]}
            goals: 目标字典 {robot_id: goal}
        """
        self.robot_paths = solution
        self.robot_goals = goals
        
        # 初始化机器人位置
        self.robot_positions = {}
        for robot_id, path in solution.items():
            if path:
                self.robot_positions[robot_id] = path[0]
        
        # 重置仿真状态
        self.current_timestep = 0
        self.completed_robots = set()
        self.conflicts = []
        
        # 计算统计信息
        self.sum_of_costs = sum(len(path) for path in solution.values())
        self.makespan = max(len(path) for path in solution.values()) if solution else 0
    
    def step(self):
        """
        执行一个时间步
        
        Returns:
            bool: 是否所有机器人都完成任务
        """
        self.current_timestep += 1
        
        # 检查是否所有机器人都完成任务
        if len(self.completed_robots) == len(self.robot_paths):
            return True
        
        # 更新机器人位置
        new_positions = {}
        for robot_id, path in self.robot_paths.items():
            if robot_id in self.completed_robots:
                # 已完成任务的机器人保持原位
                new_positions[robot_id] = self.robot_positions[robot_id]
                continue
            
            # 获取当前时间步的位置
            if self.current_timestep < len(path):
                new_positions[robot_id] = path[self.current_timestep]
            else:
                # 如果路径已经执行完，保持在最后一个位置
                new_positions[robot_id] = path[-1]
            
            # 检查是否到达目标
            if robot_id in self.robot_goals and new_positions[robot_id] == self.robot_goals[robot_id]:
                print(f"机器人 {robot_id} 在时间步 {self.current_timestep} 完成任务")
                self.completed_robots.add(robot_id)
        
        # 检测冲突
        self._detect_conflicts(new_positions)
        
        # 更新位置
        self.robot_positions = new_positions
        
        return len(self.completed_robots) == len(self.robot_paths)
    
    def _detect_conflicts(self, new_positions: Dict[int, Tuple[int, int]]):
        """
        检测冲突
        
        Args:
            new_positions: 新的机器人位置 {robot_id: position}
        """
        robot_ids = list(new_positions.keys())
        
        # 检查顶点冲突
        for i in range(len(robot_ids)):
            for j in range(i + 1, len(robot_ids)):
                robot_i = robot_ids[i]
                robot_j = robot_ids[j]
                
                # 如果两个机器人在同一位置
                if new_positions[robot_i] == new_positions[robot_j]:
                    self.conflicts.append((robot_i, robot_j, new_positions[robot_i], self.current_timestep))
                    print(f"冲突: 机器人 {robot_i} 和 {robot_j} 在时间步 {self.current_timestep} 位置 {new_positions[robot_i]} 发生冲突")
        
        # 检查边冲突
        for i in range(len(robot_ids)):
            for j in range(i + 1, len(robot_ids)):
                robot_i = robot_ids[i]
                robot_j = robot_ids[j]
                
                # 如果两个机器人交换位置
                if (robot_i in self.robot_positions and robot_j in self.robot_positions and
                    new_positions[robot_i] == self.robot_positions[robot_j] and
                    new_positions[robot_j] == self.robot_positions[robot_i]):
                    self.conflicts.append((robot_i, robot_j, (self.robot_positions[robot_i], self.robot_positions[robot_j]), self.current_timestep))
                    print(f"边冲突: 机器人 {robot_i} 和 {robot_j} 在时间步 {self.current_timestep} 交换位置")
    
    def run(self, max_timesteps: int = None):
        """
        运行仿真
        
        Args:
            max_timesteps: 最大时间步数，如果为None则运行到所有机器人完成任务
        """
        if max_timesteps is None:
            max_timesteps = 100  # 设置最大时间步为固定值100步
        
        if self.visualize:
            self._visualize_simulation(max_timesteps)
        else:
            for _ in range(max_timesteps):
                done = self.step()
                if done:
                    break
                time.sleep(self.step_delay)
    
    def _visualize_simulation(self, max_timesteps: int):
        """
        可视化仿真过程
        
        Args:
            max_timesteps: 最大时间步数
        """
        fig, ax = plt.subplots(figsize=(12, 12))  # 增加图像大小以便更好地观察
        
        # 绘制网格
        for i in range(self.grid_map.width + 1):
            ax.axvline(i, color='gray', linestyle='-', alpha=0.5)
        for i in range(self.grid_map.height + 1):
            ax.axhline(i, color='gray', linestyle='-', alpha=0.5)
        
        # 绘制障碍物
        for y in range(self.grid_map.height):
            for x in range(self.grid_map.width):
                if self.grid_map.grid[y, x] == 1:
                    ax.add_patch(plt.Rectangle((x, y), 1, 1, color='black'))
        
        # 使用更鲜明的颜色
        robot_colors = ['red', 'blue', 'green', 'purple', 'orange', 'cyan', 'magenta', 'yellow', 'lime', 'pink']
        
        # 绘制目标位置
        for robot_id, goal in self.robot_goals.items():
            color = robot_colors[robot_id % len(robot_colors)]
            ax.add_patch(plt.Rectangle((goal[0], goal[1]), 1, 1, color=color, alpha=0.3))
            ax.text(goal[0] + 0.5, goal[1] + 0.5, f'G{robot_id}', ha='center', va='center', fontsize=12, fontweight='bold')
        
        # 机器人标记
        robot_markers = {}
        robot_texts = {}
        for robot_id in self.robot_paths.keys():
            color = robot_colors[robot_id % len(robot_colors)]
            marker = ax.plot([], [], 'o', color=color, markersize=18)[0]  # 增大标记尺寸
            text = ax.text(0, 0, str(robot_id), color='white', ha='center', va='center', fontsize=10, fontweight='bold')
            robot_markers[robot_id] = marker
            robot_texts[robot_id] = text
        
        # 时间步文本和完成状态文本
        timestep_text = ax.text(0.02, 0.98, '', transform=ax.transAxes, ha='left', va='top', fontsize=12)
        status_text = ax.text(0.02, 0.94, '', transform=ax.transAxes, ha='left', va='top', fontsize=12)
        
        # 设置坐标轴
        ax.set_xlim(0, self.grid_map.width)
        ax.set_ylim(0, self.grid_map.height)
        ax.set_aspect('equal')
        ax.invert_yaxis()  # 使y轴原点在左上角
        
        plt.title('多机器人路径规划仿真 - 仓储环境', fontsize=14)
        
        def init():
            for marker in robot_markers.values():
                marker.set_data([], [])
            for text in robot_texts.values():
                text.set_position((0, 0))
                text.set_text('')
            timestep_text.set_text('')
            status_text.set_text('')
            return list(robot_markers.values()) + list(robot_texts.values()) + [timestep_text, status_text]
        
        def update(frame):
            # 执行一步仿真
            if frame > 0:
                done = self.step()
                if done and frame < max_timesteps - 1:
                    # 如果所有机器人都完成任务，提前结束动画
                    ani.event_source.stop()
            
            # 更新机器人位置
            for robot_id, marker in robot_markers.items():
                if robot_id in self.robot_positions:
                    pos = self.robot_positions[robot_id]
                    marker.set_data([pos[0] + 0.5], [pos[1] + 0.5])
                    robot_texts[robot_id].set_position((pos[0] + 0.5, pos[1] + 0.5))
                    robot_texts[robot_id].set_text(str(robot_id))
                    
                    # 如果机器人已完成任务，添加特殊标记
                    if robot_id in self.completed_robots:
                        marker.set_marker('*')  # 使用星形标记表示完成任务
                        marker.set_markersize(22)  # 增大标记尺寸
            
            # 更新时间步文本
            timestep_text.set_text(f'时间步: {self.current_timestep}')
            
            # 更新状态文本
            status_text.set_text(f'完成: {len(self.completed_robots)}/{len(self.robot_paths)} | 冲突: {len(self.conflicts)}')
            
            return list(robot_markers.values()) + list(robot_texts.values()) + [timestep_text, status_text]
        
        ani = FuncAnimation(fig, update, frames=max_timesteps, init_func=init, blit=True, interval=self.step_delay*1000)
        plt.tight_layout()
        plt.show()
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取仿真统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        # 重新检查一次完成情况，确保统计数据准确
        for robot_id, pos in self.robot_positions.items():
            if robot_id in self.robot_goals and pos == self.robot_goals[robot_id]:
                self.completed_robots.add(robot_id)
        
        return {
            'makespan': self.makespan,
            'sum_of_costs': self.sum_of_costs,
            'conflicts': len(self.conflicts),
            'completed_robots': len(self.completed_robots),
            'total_robots': len(self.robot_paths)
        } 